
#pragma once

#include <tmmintrin.h>
#include <cassert>

//////////////////////////////////////////////////////////////////////////
// Matrices, vectors
// column major storage (same as in default glsl)
// no classy tricks here, this should be lean and mean (in order to work as a buffer)
class vec4 {
    __m128 m_elems;

    friend class mat4;
public:
    vec4() { m_elems=_mm_setzero_ps(); }
    vec4(const vec4& src) { m_elems=src.m_elems; }
    vec4(float x, float y, float z, float w) { m_elems=_mm_setr_ps(x, y, z, w); }
    vec4(const __m128& src) { m_elems=src; }
    vec4& operator=(const vec4& src) { m_elems=src.m_elems; return *this; }
    vec4& operator=(const __m128& src) { m_elems=src; return *this; }
    float& operator[](int idx) { assert(idx<4); return m_elems.m128_f32[idx]; }
    operator float*() { return (float*)(m_elems.m128_f32); }
    operator __m128() const { return m_elems; }
    vec4& operator-=(const vec4& rhs) { m_elems=_mm_sub_ps(m_elems, rhs.m_elems); return *this; }
    vec4 operator-(const vec4& rhs) const { vec4 tmp(*this); return tmp-=rhs; }
    vec4& operator+=(const vec4& rhs) { m_elems=_mm_add_ps(m_elems, rhs.m_elems); return *this; }
    vec4 operator+(const vec4& rhs) const { vec4 tmp(*this); return tmp+=rhs; }
    vec4& Normalize()
    {
        const float a=3.0f,b=0.5f;
        __m128 tmp=_mm_mul_ps(m_elems, m_elems);
        tmp=_mm_hadd_ps(tmp, tmp);
        tmp=_mm_hadd_ps(tmp, tmp);

        __m128 tmp1=_mm_rsqrt_ss(tmp);
        __m128 tmp2=_mm_mul_ss(tmp1, tmp1);
        tmp1=_mm_mul_ss(tmp1, _mm_set_ss(b));
        tmp2=_mm_mul_ss(tmp2, tmp);
        tmp2=_mm_sub_ss(_mm_set_ss(a), tmp2);
        tmp=_mm_mul_ss(tmp1, tmp2);

        tmp=_mm_shuffle_ps(tmp, tmp, _MM_SHUFFLE(0, 0, 0, 0));
        m_elems=_mm_mul_ps(m_elems, tmp);
        return *this;
    }
    float Norm3() const
    {
        __m128 tmp=_mm_castsi128_ps(_mm_slli_si128(_mm_castps_si128(m_elems), 4));
        tmp=_mm_mul_ps(tmp, tmp);
        tmp=_mm_hadd_ps(tmp, tmp);
        tmp=_mm_hadd_ps(tmp, tmp);
        tmp=_mm_sqrt_ss(tmp);
        return _mm_cvtss_f32(tmp);
    }

    // Quaternions are represented as: vec[0]*i+vec[1]*j+vec[2]*k+vec[3]
    // Quaternion multiplication (note the reverse order, which the naming suggests)
    vec4& QMulRight(const vec4& lhs)
    {
        // Nontrivial stuff is in comments
        __m128 r=_mm_mul_ps(m_elems, lhs.m_elems);
        r=_mm_hadd_ps(r, r);
        r=_mm_hadd_ps(r, r);
        r=_mm_castsi128_ps(_mm_slli_si128(_mm_castps_si128(r), 12));    // Result is intentionally r1*r2 + <v1,v2>

        __m128 r1=_mm_shuffle_ps(lhs.m_elems, lhs.m_elems, _MM_SHUFFLE(3, 3, 3, 3));
        __m128 r2=_mm_shuffle_ps(m_elems, m_elems, _MM_SHUFFLE(3, 3, 3, 3));
        r1=_mm_mul_ps(r1, m_elems);
        r2=_mm_mul_ps(r2, lhs.m_elems);
        r1=_mm_add_ps(r1, r2);  // After addition the fourth coordinate is 2*r1*r2

        __m128 tmp1=_mm_mul_ps(_mm_shuffle_ps(lhs.m_elems, lhs.m_elems, _MM_SHUFFLE(3, 0, 2, 1)),
                               _mm_shuffle_ps(m_elems, m_elems, _MM_SHUFFLE(3, 1, 0, 2)));
        __m128 tmp2=_mm_mul_ps(_mm_shuffle_ps(lhs.m_elems, lhs.m_elems, _MM_SHUFFLE(3, 1, 0, 2)),
                               _mm_shuffle_ps(m_elems, m_elems, _MM_SHUFFLE(3, 0, 2, 1)));
        tmp1=_mm_sub_ps(tmp1, tmp2);    // After subtraction the fourth coordinate is zero
        tmp1=_mm_add_ps(tmp1, r1);      // Now the fourth coordinate is 2*r1*r2
        m_elems=_mm_sub_ps(tmp1, r);    // The fourth coordinate is: 2*r1*r2 - ( r1*r2 + <v1,v2> ) == r1*r2 - <v1,v2>

        return *this;
    }

    // back conjugation of vector with quaternion: q^* * v * q
    vec4 QBackConj(const vec4& q) const
    {
        __m128 qrot=q.m_elems;
        __m128 tmp1=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(m_elems, m_elems, _MM_SHUFFLE(3, 1, 0, 2)));
        __m128 tmp2=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 1, 0, 2)),
            _mm_shuffle_ps(m_elems, m_elems, _MM_SHUFFLE(3, 0, 2, 1)));
        tmp1=_mm_sub_ps(tmp2, tmp1);
        tmp2=_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3,3,3,3));
        tmp2=_mm_mul_ps(tmp2, m_elems);
        tmp1=_mm_add_ps(tmp1, tmp2);
        tmp2=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(tmp1, tmp1, _MM_SHUFFLE(3, 1, 0, 2)));
        tmp1=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 1, 0, 2)),
            _mm_shuffle_ps(tmp1, tmp1, _MM_SHUFFLE(3, 0, 2, 1)));
        tmp2=_mm_sub_ps(tmp2, tmp1);
        tmp2=_mm_add_ps(tmp2, tmp2);
        return vec4(_mm_sub_ps(m_elems, tmp2));
    }

    // back conjugation of vector with quaternion: q * v * q^*
    vec4 QConj(const vec4& q) const
    {
        __m128 qrot=q.m_elems;
        __m128 tmp1=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(m_elems, m_elems, _MM_SHUFFLE(3, 1, 0, 2)));
        __m128 tmp2=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 1, 0, 2)),
            _mm_shuffle_ps(m_elems, m_elems, _MM_SHUFFLE(3, 0, 2, 1)));
        tmp1=_mm_sub_ps(tmp1, tmp2);
        tmp2=_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3,3,3,3));
        tmp2=_mm_mul_ps(tmp2, m_elems);
        tmp1=_mm_add_ps(tmp1, tmp2);
        tmp2=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(tmp1, tmp1, _MM_SHUFFLE(3, 1, 0, 2)));
        tmp1=_mm_mul_ps(_mm_shuffle_ps(qrot, qrot, _MM_SHUFFLE(3, 1, 0, 2)),
            _mm_shuffle_ps(tmp1, tmp1, _MM_SHUFFLE(3, 0, 2, 1)));
        tmp2=_mm_sub_ps(tmp2, tmp1);
        tmp2=_mm_add_ps(tmp2, tmp2);
        return vec4(_mm_add_ps(m_elems, tmp2));
    }
};

class mat4 {
    __m128 m_cols[4];

public:
    mat4()
    {
        m_cols[0]=_mm_setr_ps(1,0,0,0);
        m_cols[1]=_mm_setr_ps(0,1,0,0);
        m_cols[2]=_mm_setr_ps(0,0,1,0);
        m_cols[3]=_mm_setr_ps(0,0,0,1);
    }
    mat4(float* elems) { for(int i=0;i<4;i++) m_cols[i]=_mm_loadu_ps(elems+4*i); }
    mat4(const mat4& src) { for(int i=0;i<4;i++) m_cols[i]=src.m_cols[i]; }
    mat4& operator=(const mat4& src) { for(int i=0;i<4;i++) m_cols[i]=src.m_cols[i]; return *this; }
    mat4 operator*(const mat4& m) const { mat4 tmp(m); return MulRight(tmp); }
    vec4& MulRight(vec4& v) const
    {
        __m128 tmp[4];

        tmp[0]=v.m_elems;
        tmp[1]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(1,1,1,1));
        tmp[2]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(2,2,2,2));
        tmp[3]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(3,3,3,3));
        tmp[0]=_mm_shuffle_ps(tmp[0], tmp[0], _MM_SHUFFLE(0,0,0,0));

        tmp[0]=_mm_mul_ps(tmp[0], m_cols[0]);
        tmp[1]=_mm_mul_ps(tmp[1], m_cols[1]);
        tmp[2]=_mm_mul_ps(tmp[2], m_cols[2]);
        tmp[3]=_mm_mul_ps(tmp[3], m_cols[3]);

        tmp[0]=_mm_add_ps(tmp[0], tmp[1]);
        tmp[2]=_mm_add_ps(tmp[2], tmp[3]);
        v.m_elems=_mm_add_ps(tmp[0], tmp[2]);
        return v;
    }
    vec4 operator*(const vec4& rhs) const
    {
        vec4 tmp(rhs);
        return MulRight(tmp);
    }
    mat4& MulRight(mat4& m) const
    {
        for(int i=0;i<4;i++) {
            __m128 tmp[4];
            __m128 right=m.m_cols[i];

            tmp[0]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(0, 0, 0, 0));
            tmp[0]=_mm_mul_ps(tmp[0], m_cols[0]);
            tmp[1]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(1, 1, 1, 1));
            tmp[1]=_mm_mul_ps(tmp[1], m_cols[1]);
            tmp[2]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(2, 2, 2, 2));
            tmp[2]=_mm_mul_ps(tmp[2], m_cols[2]);
            tmp[3]=_mm_shuffle_ps(right, right, _MM_SHUFFLE(3, 3, 3, 3));
            tmp[3]=_mm_mul_ps(tmp[3], m_cols[3]);

            tmp[0]=_mm_add_ps(tmp[0], tmp[1]);
            tmp[2]=_mm_add_ps(tmp[2], tmp[3]);
            m.m_cols[i]=_mm_add_ps(tmp[0], tmp[2]);
        }

        return m;
    }
    mat4& operator*=(const mat4& m)
    {
        mat4 tmp(m);
        MulRight(tmp);
        m_cols[0]=tmp.m_cols[0];
        m_cols[1]=tmp.m_cols[1];
        m_cols[2]=tmp.m_cols[2];
        m_cols[3]=tmp.m_cols[3];

        return *this;
    }
    mat4& operator*=(float a)
    {
        __m128 tmp=_mm_set1_ps(a);
        m_cols[0]=_mm_mul_ps(m_cols[0], tmp);
        m_cols[1]=_mm_mul_ps(m_cols[1], tmp);
        m_cols[2]=_mm_mul_ps(m_cols[2], tmp);
        m_cols[3]=_mm_mul_ps(m_cols[3], tmp);
        return *this;
    }
    mat4& operator+=(mat4& m)
    {
        m_cols[0]=_mm_add_ps(m_cols[0], m.m_cols[0]);
        m_cols[1]=_mm_add_ps(m_cols[1], m.m_cols[1]);
        m_cols[2]=_mm_add_ps(m_cols[2], m.m_cols[2]);
        m_cols[3]=_mm_add_ps(m_cols[3], m.m_cols[3]);
        return *this;
    }
    mat4& operator-=(mat4& m)
    {
        m_cols[0]=_mm_sub_ps(m_cols[0], m.m_cols[0]);
        m_cols[1]=_mm_sub_ps(m_cols[1], m.m_cols[1]);
        m_cols[2]=_mm_sub_ps(m_cols[2], m.m_cols[2]);
        m_cols[3]=_mm_sub_ps(m_cols[3], m.m_cols[3]);
        return *this;
    }
    float* NormalMatrix(float* dst) const   // Adjugate adjoint 4x3
    {
        __m128 u=m_cols[1];
        __m128 v=m_cols[2];

        __m128 res=_mm_mul_ps(_mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 1, 0, 2)));
        __m128 res1=_mm_mul_ps(_mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 1, 0, 2)));
        *(__m128*)(dst) = _mm_sub_ps(res, res1);

        u=v;
        v=m_cols[0];
        res=_mm_mul_ps(_mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 1, 0, 2)));
        res1=_mm_mul_ps(_mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 1, 0, 2)));
        *(__m128*)(dst+4) = _mm_sub_ps(res, res1);

        u=v;
        v=m_cols[1];
        res=_mm_mul_ps(_mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 1, 0, 2)));
        res1=_mm_mul_ps(_mm_shuffle_ps(v, v, _MM_SHUFFLE(3, 0, 2, 1)),
            _mm_shuffle_ps(u, u, _MM_SHUFFLE(3, 1, 0, 2)));
        *(__m128*)(dst+8) = _mm_sub_ps(res, res1);

        return dst;
    }
    operator float*() { return (float*)m_cols; }
    float& operator()(int i, int j) { return m_cols[j].m128_f32[i]; }
    __m128& operator()(int j) { return m_cols[j]; }
};

mat4 RotationMatrixFromQ(const vec4& q);
mat4 RotateAroundAxis(float x, float y, float z, float angle);
